#include "opticalflowsegmenter.h"

OpticalFlowSegmenter::OpticalFlowSegmenter(std::vector<CGMesh*>& meshlist_par, int active_par, GLCamera& camera_par, GLWindow& window_par, QWidget* parent)
	: QWidget(parent), meshlist(meshlist_par), camera(camera_par), window(window_par)
{
	this->setMinimumSize(300,100);
	this->setMaximumSize(500,500);
	this->setWindowTitle("Optical Flow Segmentation");
	this->setWindowIcon(QIcon(tr("segm.ico")));

	//set members
	active = active_par;

	srand(time(NULL));

	//create layout
	QGridLayout* layout = new QGridLayout();

	QPushButton* colorButton = new QPushButton("Colorize mesh", this);
	QPushButton* snapshotButton = new QPushButton("Take snapshot", this);
	QPushButton* flowButton = new QPushButton("Compute Flow", this);

	this->setFocusPolicy(Qt::NoFocus);

	layout->addWidget(colorButton, 0, 0, 1, 3);
	layout->addWidget(snapshotButton, 1, 0, 1, 3);
	layout->addWidget(flowButton, 2, 0, 1, 3);

	this->setLayout(layout);

	//connect signals
	connect( colorButton,	 SIGNAL(clicked()), this, SLOT(colorizeMeshFaces())   );
	connect( snapshotButton, SIGNAL(clicked()), this, SLOT(takeSnapshot())		);
	connect( flowButton,	 SIGNAL(clicked()), this, SLOT(computeFlow())		);
}

OpticalFlowSegmenter::~OpticalFlowSegmenter()
{

}

void OpticalFlowSegmenter::colorizeMeshFaces() {
	qDebug() << "Coloring...";

	//create a set of colors, at least one for each face
	std::set< vcg::Color4b > colorSet;

	for ( int i = 0; i < meshlist[active]->face.size(); i ++ ) {
		
		//pick random color
		vcg::Color4b color;
		do {
			color = vcg::Color4b ( (unsigned char) rand()%256, (unsigned char) rand()%256, (unsigned char) rand()%256, (unsigned char) 255 );
		}
		while ( colorSet.find(color) != colorSet.end() ||
                        color == vcg::Color4b(255, 0, 255, 255) );

		//add color to set
		colorSet.insert(color);

                if ( color == vcg::Color4b(255,255,255,255) )
                    qDebug() << "YES!";
		
		//color faces
		for ( int j = 0; j < meshlist.size(); j ++ )
			meshlist[j]->face[i].C() = color;
	}

	colorSet.clear();

	emit updateView();
}

void OpticalFlowSegmenter::takeSnapshot() {
	qDebug() << "Shapshooting...";

        //remove shots
        shots.clear();

	//for each mesh of meshlist take a screenshot
	//first hide ALL the meshes!
	for ( int i = 0; i < meshlist.size(); i ++ )
		emit hideMesh(i);

	for ( int i = 0; i < meshlist.size(); i ++ ) {
		//for each mesh...

		//show it and update view
		emit showMesh(i);
		emit updateView();

		//get screenshot
		//QString filename = QFileDialog::getSaveFileName( this, "Save shapshot", "", "Image files (*.png *.jpg *.bmp);;png (*.png);;jpg (*.jpg);;bmp (*.bmp)");

                QImage* snapshot = new QImage( window.width(), window.height(), QImage::Format_ARGB32 );

                //read data from framebuffer
                uchar* red = (uchar*) calloc (snapshot->width() * snapshot->height(), sizeof(uchar));
                uchar* green = (uchar*) calloc (snapshot->width() * snapshot->height(), sizeof(uchar));
                uchar* blue = (uchar*) calloc (snapshot->width() * snapshot->height(), sizeof(uchar));
                glReadPixels( (GLint) 0, (GLint) 0, (GLint) snapshot->width(), (GLint) snapshot->height(), GL_RED, GL_UNSIGNED_BYTE, red );
                glReadPixels( (GLint) 0, (GLint) 0, (GLint) snapshot->width(), (GLint) snapshot->height(), GL_GREEN, GL_UNSIGNED_BYTE, green );
                glReadPixels( (GLint) 0, (GLint) 0, (GLint) snapshot->width(), (GLint) snapshot->height(), GL_BLUE, GL_UNSIGNED_BYTE, blue );

                //save data in QImage
                for(int x = 0; x < snapshot->width(); x++)   {      //col
                   for(int y = 0; y < snapshot->height(); y++)  {   //row
                       //save image
                       snapshot->setPixel(x, snapshot->height()-y-1, qRgb(red[y*snapshot->width()+x], green[y*snapshot->width()+x], blue[y*snapshot->width()+x]));
                   }
                }

                QString filename = QFileDialog::getSaveFileName( this, "Save shapshot", "", "Image files (*.png *.jpg *.bmp);;png (*.png);;jpg (*.jpg);;bmp (*.bmp)");
                if ( !filename.isNull() ) snapshot->save(filename);

		//store in vector
		shots.push_back(snapshot);

		//hide the mesh
		emit hideMesh(i);
	}


	//last, show ALL the meshes!
	for ( int i = 0; i < meshlist.size(); i ++ )
		emit showMesh(i);

}

void OpticalFlowSegmenter::computeFlow() {
	qDebug() << "Flowing...";


        //DEBUG
        /*shots.clear();
        for ( int i = 0; i < 2; i ++ ) {
            QString filename = QFileDialog::getOpenFileName( this, "Save shapshot", "", "Image files (*.png *.jpg *.bmp);;png (*.png);;jpg (*.jpg);;bmp (*.bmp)");
            if ( !filename.isNull() )
                shots.push_back(new QImage(filename));
        }*/


	//not enough images
	if ( shots.size() < 2 ) return;

	//first example works only on the first two images

	//N maps, one for each image
        //each maps contains a color and a list of clusters of coordinates (each pixel has that specific color)
        std::vector< std::map< vcg::Color4b, std::vector<PixelCluster> > > colorClusters( shots.size() );


	//raster of each image
	for ( int im = 0; im < 2; im ++ ) {
		//now working on colorClusters[im]
                //for each element of the image
		for ( int x = 0; x < shots[im]->width(); x ++ ) {
			for ( int y = 0; y < shots[im]->height(); y ++ ) {
				//get current pixel
				vcg::Color4b current( (unsigned char) qRed(shots[im]->pixel(x, y)), 
									  (unsigned char) qGreen(shots[im]->pixel(x, y)), 
									  (unsigned char) qBlue(shots[im]->pixel(x, y)), 
									  (unsigned char) 255);


				//is it background? descard it!
				if ( current == vcg::Color4b( (unsigned char) 255, (unsigned char) 0, (unsigned char) 255, (unsigned char) 255 ) )
					continue;

				//not background? good!

				//let see if it exist in our color cluster
                                /*if ( colorClusters[im].find( current ) != colorClusters[im].end() ) {
					//it's already in the list; add the pair (x,y) to the corresponding cluster
					colorClusters[im][current].insert( ImageCoordinate(x,y) );
				}
				else {
					//not present: create a new cluster
					PixelCluster dummy; dummy.insert( ImageCoordinate(x,y) );
					//insert dummy cluster in cluster list
					colorClusters[im][current] = dummy;
                                }*/


                                if ( colorClusters[im].find( current ) != colorClusters[im].end() ) {
                                        //it's already in the list; search if one of the clusters is adjacent to the (x,y) coordinates
                                    bool found = false;
                                    for ( int i = 0; i < colorClusters[im][current].size(); i ++ )
                                        if ( isAdjacent( colorClusters[im][current][i], ImageCoordinate(x,y) ) ) {
                                            colorClusters[im][current][i].insert( ImageCoordinate(x,y) );
                                            found = true;
                                            break;  //no more iteration
                                        }

                                    if ( !found ) {
                                        //not adjacent clusters: create a new cluster
                                        PixelCluster dummy; dummy.insert( ImageCoordinate(x,y) );
                                        //insert dummy cluster in cluster list
                                        colorClusters[im][current].push_back(dummy);
                                    }

                                }
                                else {
                                        //not present: create a new cluster
                                        PixelCluster dummy; dummy.insert( ImageCoordinate(x,y) );
                                        //insert dummy cluster in cluster list
                                        colorClusters[im][current].push_back(dummy);
                                }
			}
		}
	}

	//all clusters have been created
	
	int matching = 0;

	//test with two maps
        std::map< vcg::Color4b, std::vector<PixelCluster> >::iterator map0 = colorClusters[0].begin();
        std::map< vcg::Color4b, std::vector<PixelCluster> >::iterator map1 = colorClusters[1].begin();

	std::vector< std::pair<vcg::Point2f, PixelCluster> > flow;

	float max_mod = 0;
	float min_mod = 1000;



	while ( map0 != colorClusters[0].end() && map1 != colorClusters[1].end() ) {

		if ( (*map0).first == (*map1).first ) {
			//same color: we have a matching!
			matching++;

                        std::vector<PixelCluster> tmp0 = (*map0).second;
                        std::vector<PixelCluster> tmp1 = (*map1).second;

                        /*
                        //get the barycenters
                        std::vector< vcg::Point2f > bars0( (*map0).second.size() );
                        std::vector< vcg::Point2f > bars1( (*map1).second.size() );

                        ///DA QUI!!!


                        vcg::Point2f bar0 = getBarycenter( (*map0).second[0] );
                        vcg::Point2f bar1 = getBarycenter( (*map1).second[0] );
			//compute flow vector and store with the elements of first queue
                        flow.push_back( std::pair<vcg::Point2f, PixelCluster> ( bar0-bar1, (*map0).second[0] ) );
                        max_mod = std::max( (bar0-bar1).Norm(), max_mod );
                        min_mod = std::min( (bar0-bar1).Norm(), min_mod );
                        */


                        //get the barycenters
                        std::vector< vcg::Point2f > bars0( (*map0).second.size() );
                        std::vector< vcg::Point2f > bars1( (*map1).second.size() );

                        for ( int i = 0; i < (*map0).second.size(); i ++ )
                            bars0[i] = getBarycenter( (*map0).second[i] );

                        for ( int i = 0; i < (*map1).second.size(); i ++ )
                            bars1[i] = getBarycenter( (*map1).second[i] );

                        for ( int i = 0; i < bars0.size(); i ++ ) {
                            //for each cluster of bars0 find the nearest cluster bars1
                            int nearest = 0;
                            float mindist = (bars0[i]-bars1[nearest]).Norm();
                            for ( int j = 1; j < bars1.size(); j ++ ) {
                                if ( (bars0[i]-bars1[j]).Norm() < mindist  ) {
                                    nearest = j;
                                    mindist = (bars0[i]-bars1[nearest]).Norm();
                                }
                            }

                            //compute flow vector and update max and min mod
                            flow.push_back( std::pair<vcg::Point2f, PixelCluster> ( bars0[i]-bars1[nearest], (*map0).second[i] ) );
                            max_mod = std::max( (bars0[i]-bars1[nearest]).Norm(), max_mod );
                            min_mod = std::min( (bars0[i]-bars1[nearest]).Norm(), min_mod );
                        }

			//update iterators
			map0++; map1++;
			continue;
		}

		if ( (*map0).first < (*map1).first ) {
			//color of map0 is lower -> move the iterator
			map0++;
			continue;
		}

		if ( (*map0).first > (*map1).first ) {
			//color of map1 is lower -> move the iterator
			map1++;
			continue;
		}
	
	}

	int k = 0;

	qDebug() << "Matching: " << matching;
	qDebug() << "Map 0: " << colorClusters[0].size();
	qDebug() << "Map 1: " << colorClusters[1].size();


        //flow cleaning...


	//save flow on image (DEBUG)
	QImage flowimage( shots[0]->width(), shots[0]->height(), QImage::Format_ARGB32 );
        flowimage.fill( QColor(255, 0, 255).rgba() );

	//for each vector...
	for ( int i = 0; i < flow.size(); i ++ ) {
		//get the flow vector
		vcg::Point2f vector = flow[i].first;
		QColor color(0, 0, 0);
		float x = vector.X();
		float y = vector.Y();

		int h, s, v;

		if ( vector.Norm() == 0.0 ) {
			//null vector -> white
			color.setRgb(255, 255, 255, 255);
		}
		else {
			//compute HSV

			float a = 180 * atan(y/x) / M_PI;
			if ( x > 0.0f && y == 0.0f ) h = 0;
			if ( x > 0.0f && y > 0.0f  ) h = a;
			if ( x == 0.0f && y > 0.0f ) h = 90;
			if ( x < 0.0f && y > 0.0f  ) h = 180 + a;
			if ( x < 0.0f && y == 0.0f ) h = 180;
			if ( x < 0.0f && y < 0.0f  ) h = 180 + a;
			if ( x == 0.0f && y < 0.0f ) h = 270;
			if ( x > 0.0f && y < 0.0f  ) h = 360 + a;
			//compute SATURATION (normalize in 0...255 range)
			if ( min_mod == max_mod ) 
				s = 128;
			else 
				s = 255*log10(1 + 9*(vector.Norm() - min_mod)/(max_mod - min_mod));
			//set VALUE to 255
			if ( s < 0 ) s = 0;
			if ( s > 255 ) s = 255;

			if ( h < 0 ) h = 0;
			v = 255;
			//errori nel valore di S dovuti a elementi non correttamente inizializzati
			color.setHsv( h, s, v );
		}

		//compute rgb value
		QRgb value = qRgb( color.toRgb().red(), color.toRgb().green(), color.toRgb().blue() );


		//flowimage.setPixel(i, j, value);

		//set each pixel of the cluster with the value color
		for ( PixelCluster::iterator it = flow[i].second.begin(); it != flow[i].second.end(); ++it ) {
			flowimage.setPixel( (*it).first, (*it).second, value );
		}

	}

	QString filename = QFileDialog::getSaveFileName( this, "Save shapshot", "", "Image files (*.png *.jpg *.bmp);;png (*.png);;jpg (*.jpg);;bmp (*.bmp)");
	if ( !filename.isNull() ) flowimage.save(filename);

	qDebug() << "Ok!";


	shots.clear();
}




vcg::Point2f OpticalFlowSegmenter::getBarycenter( PixelCluster pc ) {
	
	vcg::Point2f bar(0.0f, 0.0f);

	//accumulator
        for ( PixelCluster::iterator it = pc.begin(); it != pc.end(); ++ it ) {
		bar.X() += (*it).first;
		bar.Y() += (*it).second;
	}		

	//mean
        bar = bar / pc.size();

	//barycenter!
	return bar;

}


bool OpticalFlowSegmenter::isAdjacent( PixelCluster pc, ImageCoordinate query ) {

    int mindist = 1000;

    for ( PixelCluster::iterator it = pc.begin(); it != pc.end(); ++ it ) {
        int dist = std::max(  abs((*it).first - query.first), abs((*it).second - query.second)  );
        mindist = std::min(dist, mindist);
    }

    return mindist < 5;
}
